from typing import BinaryIO, Union

from application.prompt.prompt_template_service import PromptTemplateService
from domain.model.register import CHAT_MODEL_MANAGER
from infrastructure.constant import models, file
from infrastructure.excel import ExcelReader
from infrastructure.utils import env, text
from infrastructure.utils.file import save_text_content


class EntityStructGeneratorService:

    @staticmethod
    def generate_entity_struct(excel_file: Union[str, bytes, BinaryIO], model_name: str = models.DEFAULT_MODEL):
        # 启动观察
        env.set_langsmith()
        # 得到聊天模型
        chat_model = CHAT_MODEL_MANAGER.get_model_instance_by_name(model_name)
        # 读取excel数据
        json_array = ExcelReader.read_excel_as_json_array(excel_file)
        # 构造prompt
        prompt_value = PromptTemplateService.from_message_list_with_system_message(
            messages=[{
                'role': 'user', 'content': ExcelReader.to_json_str(json_array)
            }],
            system_prompt_content=PromptTemplateService.get_string_prompt_with_system_message_excel(version='v3'),
        )
        # 同步调用模型
        result_message = chat_model.invoke(input=prompt_value)
        content = result_message.content
        # 去除深度思考tag
        if "<think>" in content:
            content = text.remove_think_blocks(content).strip()
        # 去除go-markdown-code包裹
        if "```go" in content:
            content = text.extract_md_go_code_blocks(content).strip()
        # 保持生成内容到文件
        save_text_content("D:\\Code_Go\\tongquetai\server\\activity_vk", file.FILE_NAME_OF_ENTITY_STRUCT, content)
        return content
